#!/usr/bin/env python3

import rospy
import tf
from tf.transformations import euler_from_quaternion
import csv

class TfLogger:
    def __init__(self):
        rospy.init_node('tf_logger', anonymous=True)
        # 从参数服务器获取文件路径
        self.file_path = rospy.get_param('~file_path')
        # 创建TF监听器
        self.tf_listener = tf.TransformListener()
        # 设置循环频率
        self.rate = rospy.Rate(5)  # 5Hz
        # 过一段时间启动程序
        rospy.sleep(3)
        rospy.loginfo("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")

    def run(self):
        with open(self.file_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['timeStamp','x', 'y', 'z', 'roll', 'pitch', 'yaw'])
            while not rospy.is_shutdown():
                try:
                    # 获取map到body的变换
                    (trans, rot) = self.tf_listener.lookupTransform(
                        '/map', '/body', rospy.Time(0)
                    )
                    # 获取当前时间戳（秒）
                    timestamp = rospy.Time.now().to_sec()
                    # 将四元数转换为欧拉角（roll, pitch, yaw）
                    roll, pitch, yaw = euler_from_quaternion(rot)
                    # 写入CSV文件
                    writer.writerow([
                        timestamp, 
                        trans[0], trans[1], trans[2],  # x, y, z
                        roll, pitch, yaw              # roll, pitch, yaw
                    ])
                    csvfile.flush()  # 确保数据立即写入磁盘
                except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException) as e:
                    rospy.logwarn(f"TF变换获取失败: {str(e)}")
                self.rate.sleep()

if __name__ == '__main__':
    logger = TfLogger()
    logger.run()